package main

import (
	"fmt"
	"runtime"
	"sync"
	"time"
)

var wg sync.WaitGroup

// Task 任务接口
type Task interface {
	Execute()
}

// Pool 协程池
type Pool struct {
	TaskChannel chan Task // 任务队列
}

// NewPool 创建一个协程池
func NewPool(cap ...int) *Pool {
	// 获取 worker 数量
	var n int
	if len(cap) > 0 {
		n = cap[0]
	}
	if n == 0 {
		n = runtime.NumCPU()
	}

	p := &Pool{
		TaskChannel: make(chan Task),
	}

	// 创建指定数量 worker 从任务队列取出任务执行
	for i := 0; i < n; i++ {
		go func() {
			for task := range p.TaskChannel {
				task.Execute()
			}
		}()
	}
	return p
}

// Submit 提交任务
func (p *Pool) Submit(t Task) {
	p.TaskChannel <- t
}

// EatFood 吃饭任务
type AuditTask struct {
	wg *sync.WaitGroup
}

func (e *AuditTask) Execute() {
	defer e.wg.Done()
	fmt.Println("*****eat cost 3 seconds")
	time.Sleep(3 * time.Second)
}

func main() {
	p := NewPool(5)
	for i := 0; i < 10; i++ {
		wg.Add(1)
		task1 := &AuditTask{
			wg: &wg,
		}
		p.Submit(task1)
	}
	// 等待所有任务执行完成
	wg.Wait()
}
